import os,sys
print(os.getcwd())
sys.path.append(os.getcwd())
from skorch.core import *
from skorch.version import *
import json
import SimpleITK as sitk


def resample_image(itk_image, out_spacing=(1.0, 1.0, 1.0), is_label=False):
    original_spacing = itk_image.GetSpacing()
    original_size = itk_image.GetSize()

    out_size = [int(np.round(original_size[0]*(original_spacing[0]/out_spacing[0]))),
                int(np.round(original_size[1]*(original_spacing[1]/out_spacing[1]))),
                int(np.round(original_size[2]*(original_spacing[2]/out_spacing[2])))]

    resample = sitk.ResampleImageFilter()
    resample.SetOutputSpacing(out_spacing)
    resample.SetSize(out_size)
    resample.SetOutputDirection(itk_image.GetDirection())
    resample.SetOutputOrigin(itk_image.GetOrigin())
    resample.SetTransform(sitk.Transform())
    resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())

    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        resample.SetInterpolator(sitk.sitkBSpline)
    return resample.Execute(itk_image)


def vertebra_Trans(img):
    image_array = sitk.GetArrayFromImage(img)
    print(np.shape(image_array))
    vertebra_win = np.array([200., 800.])
    newimg = (image_array-vertebra_win[0])/(vertebra_win[1]-vertebra_win[0])
    newimg[newimg < 0] = 0
    newimg[newimg > 1] = 1
    newimg = (newimg*255).astype('uint8')
    return newimg


class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)


def ReadTxt(rootdir):
    lines = []
    with open(rootdir, 'r') as file_to_read:
        while True:
            line = file_to_read.readline()
            if not line:
                break
            line = Path(line.strip('\n'))
            lines.append(line)
    return lines


def patch_loader(image_items):
    image_patch = {}
    leng = len(image_items)
    for n in range(leng):
        sample_json_item = image_items[n]
        image_name = sample_json_item.name.split('_ctd')[0] + '.npy'
        mask_name = sample_json_item.name.split('_ctd')[0] + '_seg.npy'
        vertebrate_mask_item = sample_json_item.parent.joinpath(mask_name)
        # print('=======================================================', image_name)
        newimg_mask = (np.load(vertebrate_mask_item)).astype(np.uint8)
        # ==================sliding windows -> patch -> net================== #
        deep, height, width = np.shape(newimg_mask)         # 图像[deep, height, width]
        # print('np.shape(newimg) ', np.shape(newimg_mask))
        step = 12
        # m = 0
        patch_flag = False
        for d in range(-120, deep - step, step):
            for j in range(0, height - step, step):
                for i in range(0, width - step, step):
                    # print('d:{},i:{},j:{}'.format(d, j, i))
                    dd = d
                    if dd < 0:
                        dd = 0
                    mask_slid = newimg_mask[dd:128 + d, j:128 + j, i:128 + i]
                    sumVert = np.sum(mask_slid > 0)
                    if sumVert > 1000:
                        centre = np.mean(np.where(mask_slid != 0), axis=1)
                        patch_flag = True
                        break
                if patch_flag:
                    break
            if patch_flag:
                break
        patch_width = i
        patch_height = (j + abs(centre[1] - 64)).astype(int)
        patch_deep = d
        # ============================================================================ #
        for d in range(patch_deep, deep - step, step):
            image_patch['img' + str(n) + '-' + str(d) + '-' + str(patch_height)
                                 + '-' + str(patch_width)] = [n, d, patch_height, patch_width]
            image_patch['img' + str(n) + '/' + str(d) + '/' + str(0)
                        + '/' + str(0)] = [n, d, 0, 0]
    return image_patch


def prepare_dict():
    image_items = ReadTxt('train_items.txt')  # 图像序列
    patch_idx = patch_loader(image_items)
    with open('patch_dict.json', 'w') as load_file:  # patch序列
        json.dump(patch_idx, load_file, cls=NpEncoder)
    print(len(patch_idx))
    with open('patch_dict.json', 'r') as load_file:  # patch序列
        patch_idx = json.load(load_file)
        print(len(patch_idx))
    patch_dict = {}
    for n in range(len(patch_idx)):
        image_patch_idx = list(patch_idx.values())[n]
        sample_json_item = image_items[image_patch_idx[0]]
        image_name = sample_json_item.name.split('_ctd')[0] + '.npy'
        mask_name = sample_json_item.name.split('_ctd')[0] + '_seg.npy'
        vertebrate_img_item = sample_json_item.parent.joinpath(image_name)
        vertebrate_mask_item = sample_json_item.parent.joinpath(mask_name)
        img = np.load(vertebrate_img_item).astype(np.uint8)
        img_mask = (np.load(vertebrate_mask_item)).astype(np.uint8)
        dd = image_patch_idx[1]
        if dd < 0:
            dd = 0
        image_patch = img[dd:128 + image_patch_idx[1], image_patch_idx[2]:128 + image_patch_idx[2],
                          image_patch_idx[3]:128 + image_patch_idx[3]]
        mask_patch = img_mask[dd:128 + image_patch_idx[1], image_patch_idx[2]:128 + image_patch_idx[2],
                     image_patch_idx[3]:128 + image_patch_idx[3]]
        # ====================================================================== #
        mask_counter = Counter(img_mask.flatten())
        patch_counter = Counter(mask_patch.flatten())
        list_patch_counter = list(patch_counter)
        intact = []
        for i in range(1, len(list_patch_counter)):
            if list_patch_counter[i] != 0:
                a = patch_counter[list_patch_counter[i]]
                b = mask_counter[list_patch_counter[i]]
                if a > b * 0.98:
                    intact.append(list_patch_counter[i])
        name = str(image_patch_idx[0]) + '-' + str(image_patch_idx[1]) + '-' + str(image_patch_idx[2]) +\
                   '-' + str(image_patch_idx[3])
        # print('name:', name)
        patch_dict['img' + name] = intact
        # print(patch_dict)
        if image_patch.shape != (128, 128, 128) or image_patch.shape != (128, 128, 128):
            image_patch_pad = np.pad(image_patch, ((0, 128 - image_patch.shape[0]), (0, 128 - image_patch.shape[1]),
                                                   (0, 128 - image_patch.shape[2])), 'constant', constant_values=(0, 0))
            mask_patch_pad = np.pad(mask_patch, ((0, 128 - image_patch.shape[0]), (0, 128 - image_patch.shape[1]),
                                                 (0, 128 - image_patch.shape[2])), 'constant', constant_values=(0, 0))
            # print("image_patch_new:", np.shape(image_patch_new))
            # print("mask_patch_new:", np.shape(mask_patch_new))
            np.save('img' + name + '.npy', image_patch_pad)
            np.save('mask' + name + '.npy', mask_patch_pad)
        else:
            # print("image_patch:", np.shape(image_patch))
            # print("mask_patch:", np.shape(mask_patch))
            np.save('img' + name + '.npy', image_patch)
            np.save('mask' + name + '.npy', mask_patch)
    with open('intact_dict.json', 'w') as dump_file:
        json.dump(patch_dict, dump_file, cls=NpEncoder)


def get_tensor(sample_idx):
    image_name = sample_idx.name
    mask_name = 'mask' + sample_idx.name.split('img')[1]
    img_item = sample_idx.parent.joinpath(image_name)
    mask_item = sample_idx.parent.parent.joinpath('mask', mask_name)
    img = np.load(img_item).astype(np.uint8)  # 图像[deep, height, width]
    mask = (np.load(mask_item)).astype(np.uint8)
    img = img[..., np.newaxis]
    return img, mask


def image2tensor(im, stats=None):
    img = np.moveaxis(im / (255. if im.dtype == np.uint8 else 1), -1, 0).astype(np.float32)
    im_mean, im_std = np.mean(im), np.std(im)
    img = (img - im_mean + 1e-7) / (im_std + 1e-7)
    return torch.from_numpy(img)


def mask2tensor(mask):
    mask = np.expand_dims(mask / (255. if mask.dtype == np.uint8 else 1), 0).astype(np.float32)
    return torch.from_numpy(mask)


def image_nii2npy(dcm_file: Path, save_dir: Path):
    print('file', dcm_file)
    vertebrate_img = sitk.ReadImage(str(dcm_file))
    itk_img = resample_image(vertebrate_img, out_spacing=(1.0, 1.0, 1.0), is_label=False)
    img = vertebra_Trans(itk_img)
    print(np.shape(img))
    d_shape = 128 / img.shape[0]
    newimg = scipy.ndimage.interpolation.zoom(img, (d_shape, 1, 1), order=0, mode='nearest')
    print('newimg', np.shape(newimg))
    npy_name = dcm_file.with_suffix(".npy")
    if not save_dir.exists():
        save_dir.mkdir()
    npy_name = save_dir.joinpath(npy_name.name)
    np.save(str(npy_name), img)


if __name__ == '__main__':
    print('===start===')
    prepare_dict()
    image_nii2npy()
    # img = sitk.GetArrayFromImage(sitk.ReadImage('/home/cao/PycharmProjects/vertebra_segment/verse262.nii'))
    # img_mask = sitk.GetArrayFromImage(sitk.ReadImage('/home/cao/PycharmProjects/vertebra_segment/verse262_seg.nii'))
    # # itk_img = resample_image(img, out_spacing=(1.0, 1.0, 1.0), is_label=False)
    # # img = vertebra_Trans(itk_img)
    # patch_flag = False
    # step = 12
    # width, deep, height = np.shape(img)
    # for d in range(-120, deep - step, step):
    #     for j in range(0, height - step, step):
    #         for i in range(0, width - step, step):
    #             print('d:{},i:{},j:{}'.format(d, j, i))
    #             dd = d
    #             if dd < 0:
    #                 dd = 0
    #             mask_slid = img_mask[i:128 + i, dd:128 + d, j:128 + j]
    #             sumVert = np.sum(mask_slid > 0)
    #             if sumVert > 1000:
    #                 centre = np.mean(np.where(mask_slid != 0), axis=1)
    #                 patch_flag = True
    #                 break
    #         if patch_flag:
    #             break
    #     if patch_flag:
    #         break
    # patch_width = i
    # patch_height = (j + abs(centre[1] - 64)).astype(int)
    # patch_deep = d
    # # ============================================================================ #
    # print(patch_deep)
    # print(deep - step)
    #
    # for d in range(patch_deep, deep - step, step):
    #     print(d)
    #     patch = img[i:128 + i, d:128 + d, j:128 + j]
    #     out = sitk.GetImageFromArray(patch)
    #     sitk.WriteImage(out, 'out'+str(d) +'.nii')
    print('===done===')